import argparse
import re
from datetime import timedelta

from log_manager import LogManager


def format_srt(content, split_chars="，。,.；;"):
    """格式化SRT字幕"""
    # 分割为字幕块
    subtitle_blocks = re.split(r"\n\n+", content.strip())

    new_subtitles = []
    subtitle_index = 1
    is_first = True
    for block in subtitle_blocks:
        lines = block.split("\n")
        if len(lines) >= 3:
            time_line = lines[1]
            text = " ".join(lines[2:])

            start_time, end_time = time_line.split(" --> ")

            # 拆分字幕
            split_segments = split_subtitle(text, start_time, end_time, split_chars)

            # 生成新的字幕块
            for seg_start, seg_end, seg_text in split_segments:

                start_time = parse_time(seg_start)
                end_time = parse_time(seg_end)
                if is_first:
                    is_first = False
                    if start_time > timedelta() and start_time < timedelta(
                        hours=1
                    ):
                        temp_start = format_time(timedelta())
                        new_subtitles.append(
                            {
                                "index": subtitle_index,
                                "start_time": temp_start,
                                "end_time": seg_start,
                                "text": "---",
                            }
                        )
                        last_end_time = start_time
                        subtitle_index += 1
                else:
                    if start_time - last_end_time > timedelta(seconds=1):
                        new_subtitles.append(
                            {
                                "index": subtitle_index,
                                "start_time": format_time(last_end_time),
                                "end_time": seg_start,
                                "text": "---",
                            }
                        )

                        last_end_time = start_time
                        subtitle_index += 1
                    else:
                        new_subtitles[-1]["end_time"] = seg_start

                # new_subtitles.append(
                #     f"{subtitle_index}\n{seg_start} --> {seg_end}\n{seg_text}\n"
                # )
                new_subtitles.append(
                    {
                        "index": subtitle_index,
                        "start_time": seg_start,
                        "end_time": seg_end,
                        "text": seg_text,
                    }
                )
                last_end_time = end_time
                subtitle_index += 1
    res_subtitles = []
    for subtitle in new_subtitles:
        res_subtitles.append(
            f"{subtitle['index']}\n{subtitle['start_time']} --> "
            f"{subtitle['end_time']}\n{subtitle['text']}\n"
        )
    return "\n".join(res_subtitles)


def parse_time(time_str):
    """将SRT时间格式转换为timedelta对象"""
    hours, minutes, seconds = time_str.split(":")
    seconds, milliseconds = seconds.split(",")
    return timedelta(
        hours=int(hours),
        minutes=int(minutes),
        seconds=int(seconds),
        milliseconds=int(milliseconds),
    )


def format_time(td):
    """将timedelta对象转换为SRT时间格式"""
    total_seconds = td.total_seconds()
    hours = int(total_seconds // 3600)
    minutes = int((total_seconds % 3600) // 60)
    seconds = int(total_seconds % 60)
    milliseconds = int((total_seconds * 1000) % 1000)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"


def split_subtitle(text, start_time, end_time, split_chars):
    """按标点符号拆分字幕,并按字符数分配时间"""
    # 使用正则表达式在逗号和句号处分割,但保留引号内的完整性
    segments = []
    current_segment = ""
    in_quotes = False
    for char in text:
        if char == '"' or char == '"' or char == '"':
            in_quotes = not in_quotes
            current_segment += char
        elif char in split_chars and not in_quotes:
            if current_segment:
                segments.append(current_segment)
                current_segment = ""
        else:
            current_segment += char

    if current_segment:
        segments.append(current_segment)

    # 计算总时长和每个字符平均时长
    punctuation_time = 0.25
    total_duration = (parse_time(end_time) - parse_time(start_time)).total_seconds() - (
        len(segments) - 1
    ) * punctuation_time
    total_chars = sum(len(s) for s in segments)
    time_per_char = total_duration / total_chars if total_chars > 0 else 0
    # 为每段分配时间
    result = []
    current_time = parse_time(start_time)

    for i, segment in enumerate(segments):
        if not segment.strip():
            continue
        duration = timedelta(seconds=len(segment) * time_per_char)
        if i != len(segments) - 1:
            duration += timedelta(seconds=punctuation_time)
        end = current_time + duration
        result.append((format_time(current_time), format_time(end), segment.strip()))
        current_time = end

    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="SRT字幕格式化工具")
    parser.add_argument("input", help="输入的SRT文件路径")
    parser.add_argument(
        "--output",
        "-o",
        help="输出的SRT文件路径，默认为input_formatted.srt",
    )
    parser.add_argument(
        "--split-chars",
        help="用于分割的标点符号，默认为，。,.",
        default="，。,.：;",
    )

    args = parser.parse_args()

    # 读取输入文件
    with open(args.input, "r", encoding="utf-8") as f:
        content = f.read()

    # 格式化字幕
    formatted_content = format_srt(content, args.split_chars)

    # 确定输出路径
    output_path = args.output or args.input.rsplit(".", 1)[0] + "_formatted.srt"

    # 保存格式化后的字幕
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(formatted_content)

    LogManager.log(f"格式化完成，已保存到: {output_path}")
